# -*- coding: utf-8 -*-

import argparse
import subprocess
import math
import torch
import torch.nn as nn
from torch.autograd import Variable
import numpy as np

import float_tensor_pb2


class InvertedResidual(nn.Module):
    def __init__(self, inp, oup, stride, expand_ratio):
        super(InvertedResidual, self).__init__()
        self.stride = stride
        assert stride in [1, 2]

        self.use_res_connect = self.stride == 1 and inp == oup

        self.conv = nn.Sequential(
            # pw
            nn.Conv2d(inp, inp * expand_ratio, 1, 1, 0, bias=False),
            nn.BatchNorm2d(inp * expand_ratio),
            nn.ReLU(inplace=True),
            # dw
            nn.Conv2d(inp * expand_ratio, inp * expand_ratio, 3, stride, 1, groups=inp * expand_ratio, bias=False),
            nn.BatchNorm2d(inp * expand_ratio),
            nn.ReLU(inplace=True),
            # pw-linear
            nn.Conv2d(inp * expand_ratio, oup, 1, 1, 0, bias=False),
            nn.BatchNorm2d(oup),
        )
        self._initialize_weights()

    def _initialize_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
                m.weight.data.normal_(0, math.sqrt(2. / n))
                if m.bias is not None:
                    m.bias.data.uniform_() * 10
            elif isinstance(m, nn.BatchNorm2d):
                m.weight.data.uniform_() * 10
                m.bias.data.uniform_() * 10
            elif isinstance(m, nn.Linear):
                m.weight.data.normal_(0, 0.01)
                m.bias.data.zero_()

    def forward(self, x):
        if self.use_res_connect:
            return x + self.conv(x)
        else:
            return self.conv(x)


def save_float_tensor(float_tensor_np, save_file_path):
    assert isinstance(float_tensor_np, np.ndarray)
    assert float_tensor_np.ndim == 3

    float_tensor_pb = float_tensor_pb2.FloatTensor()
    float_tensor_pb.shape.channel = float_tensor_np.shape[0]
    float_tensor_pb.shape.height = float_tensor_np.shape[1]
    float_tensor_pb.shape.width = float_tensor_np.shape[2]
    float_tensor_pb.data.extend(float_tensor_np.flatten())

    with open(save_file_path, 'wb') as save_file:
        save_file.write(float_tensor_pb.SerializeToString())


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--in_channels", type=int)
    parser.add_argument("--out_channels", type=int)
    parser.add_argument("--stride", type=int)
    parser.add_argument("--expand_ratio", type=int)

    args = parser.parse_args()
    in_channels = args.in_channels
    out_channels = args.out_channels
    stride = args.stride
    expand_ratio = args.expand_ratio

    model = InvertedResidual(inp=in_channels,
                             oup=out_channels,
                             stride=stride,
                             expand_ratio=expand_ratio)
    model.eval()

    input_x = Variable(torch.randn(1, in_channels, 144, 160) * 10 + 1)
    output_y = model(input_x)

    print('output 10 front values: ')
    print(output_y.data.numpy()[0][0, 0, 0:10])

    input_file_path = '../../test_data/depthwise_conv_test_case_input.pb'
    output_file_path = '../../test_data/depthwise_conv_test_case_output.pb'
    onnx_file_path = '../../test_data/depthwise_conv_test_case.onnx'
    save_float_tensor(input_x.data.numpy()[0], input_file_path)
    save_float_tensor(output_y.data.numpy()[0], output_file_path)
    torch.onnx.export(model, input_x, onnx_file_path, verbose=False)

    subprocess.call(['../onnx2ncnn',
                     '../../test_data/depthwise_conv_test_case.onnx',
                     '../../test_data/depthwise_conv_test_case.param',
                     '../../test_data/depthwise_conv_test_case.bin'])
    subprocess.call(['rm', '../../test_data/depthwise_conv_test_case.onnx'])


if __name__ == "__main__":
    main()
